#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
MongoDB连接和操作
"""

import logging
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure

logger = logging.getLogger(__name__)

# 全局MongoDB客户端
mongo_client = None
db = None

def init_db(app):
    """初始化MongoDB连接"""
    global mongo_client, db

    mongo_uri = app.config['MONGO_URI']

    try:
        mongo_client = MongoClient(mongo_uri)
        # 验证连接
        mongo_client.admin.command('ping')
        logger.info("MongoDB连接成功")

        # 获取数据库
        db_name = app.config.get('MONGO_DB', 'admin')
        db = mongo_client[db_name]

        # 创建索引
        create_indexes()

        return True
    except ConnectionFailure as e:
        logger.error(f"MongoDB连接失败: {e}")
        return False

def create_indexes():
    """创建必要的索引"""
    global db
    if db is None:
        logger.error("数据库未初始化，无法创建索引")
        return

    # 会话集合索引
    db.sessions.create_index("session_id", unique=True)

    # 服务信息集合索引
    db.services.create_index([("session_id", 1), ("service_name", 1)], unique=True)

    # 请求集合索引
    db.requests.create_index("request_id", unique=True)
    db.requests.create_index([("session_id", 1), ("service_name", 1)])

    # 响应集合索引
    db.responses.create_index("request_id")
    db.responses.create_index([("session_id", 1), ("service_name", 1)])

    # 链路追踪相关索引
    # _id 字段默认就是唯一的，不需要显式指定 unique=True

    # 链路服务索引
    db.trace_services.create_index("trace_id")
    db.trace_services.create_index("service_name")

    # 链路拓扑索引
    db.trace_topology.create_index("service_id")
    db.trace_topology.create_index("parent_id")

    # 链路请求索引
    db.trace_requests.create_index("service_id")
    db.trace_requests.create_index("request_type")

    logger.info("MongoDB索引创建完成")

def get_db():
    """获取数据库连接"""
    global db
    if db is None:
        raise RuntimeError("数据库未初始化")
    return db
